import numpy as np
import matplotlib.pyplot as plt
import os

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from simulation import Simulation
from visualization import Visualization
from utils import ensure_directory_exists, save_simulation_data, print_simulation_summary, set_plot_style

def example_circle_trajectory():
    """圆形轨迹跟踪示例"""
    print("\n=== 圆形轨迹跟踪示例 ===")
    
    # 创建仿真环境
    sim = Simulation(dt=0.1, sim_time=30)
    
    # 设置LQR控制器参数（位置权重高，更注重位置跟踪精度）
    Q = np.diag([10, 10, 1, 1])
    R = np.diag([1, 1])
    sim.controller.update_parameters(Q=Q, R=R)
    
    # 重置仿真
    sim.reset(trajectory_type='circle')
    
    # 运行仿真
    sim.run()
    
    # 获取结果
    results = sim.get_results()
    errors = sim.calculate_errors()
    
    # 打印摘要
    print_simulation_summary(errors)
    
    # 可视化结果
    viz = Visualization()
    viz.plot_trajectory(results, errors)
    
    # 创建动画
    anim = viz.create_animation(results)
    
    return results, errors

def example_figure_eight_trajectory():
    """8字形轨迹跟踪示例"""
    print("\n=== 8字形轨迹跟踪示例 ===")
    
    # 创建仿真环境
    sim = Simulation(dt=0.1, sim_time=30)
    
    # 设置LQR控制器参数（平衡位置和速度跟踪）
    Q = np.diag([8, 8, 2, 2])
    R = np.diag([0.5, 0.5])
    sim.controller.update_parameters(Q=Q, R=R)
    
    # 重置仿真
    sim.reset(trajectory_type='figure_eight')
    
    # 运行仿真
    sim.run()
    
    # 获取结果
    results = sim.get_results()
    errors = sim.calculate_errors()
    
    # 打印摘要
    print_simulation_summary(errors)
    
    # 可视化结果
    viz = Visualization()
    viz.plot_trajectory(results, errors)
    
    # 创建动画
    anim = viz.create_animation(results)
    
    return results, errors

def example_square_trajectory():
    """方形轨迹跟踪示例"""
    print("\n=== 方形轨迹跟踪示例 ===")
    
    # 创建仿真环境
    sim = Simulation(dt=0.1, sim_time=40)  # 方形轨迹需要更长的仿真时间
    
    # 设置LQR控制器参数（更高的位置权重，应对方形拐角处的挑战）
    Q = np.diag([15, 15, 5, 5])
    R = np.diag([0.1, 0.1])
    sim.controller.update_parameters(Q=Q, R=R)
    
    # 重置仿真
    sim.reset(trajectory_type='square')
    
    # 运行仿真
    sim.run()
    
    # 获取结果
    results = sim.get_results()
    errors = sim.calculate_errors()
    
    # 打印摘要
    print_simulation_summary(errors)
    
    # 可视化结果
    viz = Visualization()
    viz.plot_trajectory(results, errors)
    
    # 创建动画
    anim = viz.create_animation(results)
    
    return results, errors

def example_parameter_tuning():
    """参数调优示例"""
    print("\n=== LQR参数调优示例 ===")
    
    # 创建输出目录
    ensure_directory_exists('tuning_results')
    
    # 设置绘图样式
    set_plot_style()
    
    # 不同的Q矩阵权重组合
    q_values = [
        {'name': 'high_pos_low_vel', 'values': np.diag([20, 20, 1, 1])},
        {'name': 'balanced', 'values': np.diag([10, 10, 5, 5])},
        {'name': 'high_vel_low_pos', 'values': np.diag([5, 5, 15, 15])}
    ]
    
    # 保存各组参数的RMSE结果
    rmse_results = []
    
    # 创建图形
    plt.figure(figsize=(12, 8))
    
    # 对每组参数进行测试
    for q_set in q_values:
        print(f"\n测试参数组: {q_set['name']}")
        
        # 创建仿真环境
        sim = Simulation(dt=0.1, sim_time=30)
        
        # 设置LQR控制器参数
        sim.controller.update_parameters(Q=q_set['values'])
        
        # 重置仿真
        sim.reset(trajectory_type='circle')
        
        # 运行仿真
        sim.run()
        
        # 获取结果
        results = sim.get_results()
        errors = sim.calculate_errors()
        
        # 打印摘要
        print_simulation_summary(errors)
        
        # 保存RMSE
        rmse_results.append({'name': q_set['name'], 'rmse': errors['rmse']})
        
        # 绘制轨迹
        plt.plot(results['states_history'][:, 0], results['states_history'][:, 1], 
                 label=f"{q_set['name']} (RMSE: {errors['rmse']:.4f})")
    
    # 绘制参考轨迹
    plt.plot(results['reference_trajectory'][:, 0], results['reference_trajectory'][:, 1], 
             'k--', label='参考轨迹')
    
    plt.xlabel('X位置')
    plt.ylabel('Y位置')
    plt.title('不同LQR参数下的轨迹跟踪对比')
    plt.legend()
    plt.grid(True)
    plt.axis('equal')
    
    plt.savefig('tuning_results/parameter_tuning_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 绘制RMSE柱状图
    plt.figure(figsize=(10, 6))
    names = [item['name'] for item in rmse_results]
    rmse_values = [item['rmse'] for item in rmse_results]
    
    plt.bar(names, rmse_values)
    plt.xlabel('参数组')
    plt.ylabel('RMSE')
    plt.title('不同参数组的RMSE对比')
    plt.savefig('tuning_results/rmse_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return rmse_results

if __name__ == "__main__":
    # 设置绘图样式
    set_plot_style()
    
    # 运行圆形轨迹示例
    results_circle, errors_circle = example_circle_trajectory()
    
    # 运行8字形轨迹示例
    results_figure_eight, errors_figure_eight = example_figure_eight_trajectory()
    
    # 运行方形轨迹示例
    results_square, errors_square = example_square_trajectory()
    
    # 运行参数调优示例
    # rmse_results = example_parameter_tuning()
    
    print("\n所有示例完成!") 